import unittest
import numpy as np
import math
import pickle as pkl

from causalai.models.common.prior_knowledge import PriorKnowledge
from causalai.data.time_series import TimeSeriesData
from causalai.data.transforms.time_series import StandardizeTransform
from causalai.models.time_series.base import BaseTimeSeriesAlgo

class TestBaseTimeSeriesModel(unittest.TestCase):
    def get_data(self):
        data = [[ 3.53021827e-01, -5.28816737e-01, -5.15807187e-01,
                 1.51772597e+00,  1.48093713e-01, -8.31876486e-01],
               [ 8.93720263e-01,  1.24557787e-01, -1.20677816e+00,
                -1.06292964e+00,  6.65226500e-02,  1.84621174e+00],
               [-7.38830047e-01, -2.10041348e+00,  1.46357248e+00,
                -1.08259926e+00,  5.81059475e-01,  1.31778162e+00],
               [ 1.67883629e-01, -1.11042220e+00,  1.88784771e-01,
                -2.23510169e-01,  1.59412095e+00,  5.88759083e-01],
               [ 4.31552977e-01, -1.91936941e-01,  1.18445744e+00,
                -4.13338296e-01,  1.59669837e-01,  7.50212577e-01],
               [ 4.51703853e-01,  1.52270498e+00, -4.23436014e-01,
                 3.49454441e-01,  8.57221457e-02,  2.11886605e-01],
               [ 5.57305892e-01, -1.39763207e-01,  9.68385122e-01,
                 3.21155676e-01, -7.09241999e-01,  3.22384525e-01],
               [-1.64621786e+00, -1.96939243e-01,  8.04031807e-01,
                 3.30647750e-01, -1.32842669e+00, -6.94471710e-01],
               [ 5.22677710e-02, -1.55898851e+00, -1.67367490e+00,
                 6.85118487e-01, -1.01247860e-01, -3.94841011e-01],
               [ 1.03904563e-01, -1.31514056e-02,  5.34273522e-01,
                 1.43522259e+00,  1.09033686e+00,  1.40351717e-01],
               [-1.01635017e+00, -1.25744850e+00,  5.35858808e-01,
                -9.68792403e-01, -2.03075837e-01,  6.62749719e-01],
               [ 1.68808749e+00,  8.31435528e-01, -8.05539716e-01,
                 1.29597001e-01, -1.11117625e+00, -4.05385403e-01],
               [ 9.50404673e-02,  2.22253980e+00,  8.31203877e-01,
                 9.90417711e-01,  1.68386007e+00, -2.58844073e-01],
               [ 1.71280010e+00,  5.06928175e-01, -7.85837235e-01,
                 4.84593811e-01,  4.88127832e-01,  2.63094528e-01],
               [-9.09645286e-01, -1.01832383e+00,  4.44207455e-01,
                -8.39573746e-01,  4.84397682e-01,  2.98289463e-01],
               [-8.88763783e-01,  5.69038352e-01,  1.19742569e-01,
                -9.96192318e-01, -1.50537976e-01, -1.95907221e-01],
               [ 2.62535931e-01, -8.20185139e-01,  2.06791622e+00,
                 1.51722480e+00, -6.98350521e-01, -1.81889092e+00],
               [-1.44206779e+00, -3.07179745e-01,  1.89917600e+00,
                 1.61050795e+00, -1.09687030e+00,  2.02247597e-01],
               [ 3.85292468e-01,  3.44702634e-01, -8.90358410e-01,
                 4.01325808e-01, -1.01466961e+00, -1.01727953e+00],
               [-1.73705701e-01,  1.43724338e+00,  9.22019012e-01,
                 3.57355410e-01,  4.90969611e-01, -3.41204246e-01],
               [ 1.12809194e+00,  1.61255400e-02,  1.97166100e+00,
                -1.36884192e+00,  5.14988545e-02,  1.54036769e+00],
               [ 7.96479488e-01,  6.90859646e-01, -1.42882876e+00,
                 1.02174723e+00,  8.46902772e-01,  1.77595832e-02],
               [-5.94480651e-01,  2.63689565e-01,  9.27575379e-02,
                 2.41622663e-02,  5.67483728e-01, -1.30263436e-01],
               [-9.34184132e-01,  4.75708697e-01, -7.49143616e-01,
                -1.59511498e-01, -1.25366834e-01, -1.61214234e+00],
               [ 1.63916682e+00, -9.34757827e-02, -7.64548231e-01,
                -1.93352813e+00, -2.06370288e-01,  4.92985503e-01],
               [-8.81596187e-01,  1.15326593e+00,  1.65447356e-02,
                 1.32181271e+00, -1.90683947e-01,  7.80556733e-01],
               [-1.20508432e+00, -1.51037490e+00,  8.34997493e-01,
                -2.54394773e-01,  1.04238023e-01, -7.19573511e-01],
               [-9.52057760e-01, -7.59731891e-01,  9.65447053e-01,
                 1.33409927e-01,  1.12353624e+00, -2.06905992e+00],
               [ 1.29190569e+00, -1.02940844e+00,  1.11251219e+00,
                -4.61885611e-02, -1.54386021e+00,  4.40647898e-01],
               [-4.79606337e-01, -5.96746073e-01, -2.10475190e-01,
                -1.31149960e-03,  7.57516661e-01,  8.92659678e-01],
               [-9.81657977e-01, -8.68033625e-01,  7.03108170e-01,
                -1.20590263e-01, -6.41058552e-01, -1.49760943e+00],
               [ 2.52065433e-01,  4.73335887e-01, -3.76102667e-01,
                 8.16095366e-01, -2.65814680e-01,  5.75158308e-01],
               [-1.55289010e-01, -5.80690929e-01,  3.70416974e-01,
                 1.84157581e+00, -1.35528249e-01,  1.38504919e+00],
               [-7.90814706e-01, -7.40232700e-02, -2.72863550e+00,
                 5.52574992e-01, -2.28059578e+00,  3.13962956e-01],
               [ 5.38643748e-01, -4.98024271e-01,  2.40794046e-01,
                 1.38715898e+00, -8.22379002e-01,  8.71209797e-01],
               [ 1.22789982e+00, -1.02979410e+00, -2.83280047e-01,
                -3.46186610e-01, -7.19780071e-01, -1.08770374e+00],
               [ 9.88385800e-01,  1.24274292e-01,  1.61654002e+00,
                 1.39839232e+00, -3.19708361e-01, -4.78479487e-01],
               [ 1.01569465e+00, -1.33657898e+00,  1.34096309e+00,
                 1.44265480e-01, -1.32574858e+00,  2.64110094e-01],
               [ 4.27951467e-01, -9.27460290e-01,  1.12096228e+00,
                 1.31120517e+00,  7.09957490e-01,  1.60200976e+00],
               [-1.71163972e+00, -1.52992261e+00,  3.50100363e-01,
                 8.07501374e-01, -4.01121158e-01,  9.39884927e-01],
               [ 1.73067185e+00, -5.88183652e-01,  3.65839837e-01,
                -1.91329056e+00,  5.93738079e-01, -2.59466853e-01],
               [ 4.98774170e-01, -4.09173638e-02, -3.43539997e-01,
                -3.67096453e-01,  4.22124279e-02, -2.86473196e-01],
               [ 6.67790921e-01,  1.19649385e+00,  1.45132568e+00,
                -1.11455635e+00, -3.29700059e-01,  2.07593779e+00],
               [ 1.23414821e+00, -8.70857930e-01,  1.40120382e+00,
                -1.82391970e-01,  3.15390418e-01,  1.34298107e+00],
               [ 1.54482557e+00,  1.72723698e+00, -3.06372933e-01,
                 1.58521597e-01, -1.48888931e-02, -4.66587668e-01],
               [ 5.95354497e-01,  1.29405517e-02,  8.31665519e-01,
                 3.26465767e+00, -3.65610066e-01, -7.89837346e-01],
               [ 6.53010818e-01,  4.03018845e-01, -7.47774191e-01,
                -4.51016448e-01,  5.77617161e-02, -7.13819909e-02],
               [-7.20660463e-01,  2.52575043e-01,  1.12011580e+00,
                 9.66206180e-01, -1.03610038e+00,  1.88957234e+00],
               [ 6.61175962e-02,  4.97687423e-01,  9.54207102e-01,
                 7.95365496e-01, -8.97474306e-01,  1.07871081e+00],
               [-7.65878750e-01,  1.94190481e+00,  6.10504262e-02,
                -5.36364310e-02, -1.71534578e+00,  4.90145129e-01],
               [-5.24878256e-01,  1.51407177e+00,  2.31256037e+00,
                 4.29148272e-01, -7.47685446e-01, -1.03031622e+00],
               [-5.54447285e-01, -9.71522736e-01, -1.38571252e-01,
                -4.54906052e-01, -2.47489996e-01,  8.03786273e-01],
               [-1.20866503e-01, -1.88310782e-01, -1.15393659e+00,
                 4.14253420e-01, -1.74693202e+00,  2.59346979e+00],
               [ 8.30667907e-01,  4.15857213e-01,  1.51149223e+00,
                 4.55134439e-01, -1.27253606e-01, -2.05643027e-01],
               [ 3.40250262e+00, -1.17872152e+00,  4.13063159e-01,
                 8.75318123e-01, -3.06222753e-01,  1.83581577e-01],
               [ 4.90152933e-01,  5.42331571e-01,  1.39408755e+00,
                -9.16879848e-02, -4.85421775e-02, -2.85207481e-01],
               [-1.22956864e+00, -3.34796230e-02,  1.15284724e+00,
                -1.52116689e-01,  2.56329756e-02, -3.13816047e-01],
               [ 1.39862621e+00, -8.59952013e-01,  8.92509006e-01,
                -1.03796868e+00,  5.44002925e-01,  3.07297642e+00],
               [-2.40053724e-01,  4.59687722e-01, -1.31549876e+00,
                -4.21875573e-01, -1.55970909e+00, -1.07788492e+00],
               [-8.23540075e-01,  1.06293206e+00, -4.06348504e-01,
                 1.42213825e+00, -3.34612761e-02, -1.70133270e+00],
               [ 4.76512270e-01, -1.07463664e+00, -1.05001717e+00,
                 1.08644381e+00, -1.43989830e+00, -2.73669095e-01],
               [-1.45341481e+00,  1.56196035e+00,  1.59611283e+00,
                -3.17682253e-01, -1.30297122e+00,  6.49066410e-01],
               [ 2.11941177e-01,  1.32614057e-01, -8.38195909e-01,
                 4.26507258e-01, -1.90079817e+00, -6.90448775e-02],
               [-1.99562024e-01,  3.14156028e-02,  2.05738643e+00,
                 1.50524233e+00, -1.47685171e+00,  2.20087104e+00],
               [ 3.69023191e-01,  3.08502965e-01,  8.48532621e-01,
                -4.01366200e-01, -1.73654442e+00, -1.34285831e+00],
               [ 4.59365029e-01, -2.38715534e+00, -1.48400484e+00,
                -1.87402153e-01,  4.30733686e-01, -8.36920555e-01],
               [ 1.39178877e-01,  8.60144725e-01,  3.25052579e-01,
                 3.85625728e-01, -1.81880932e+00,  1.96939825e+00],
               [-8.54424893e-02,  1.16516290e-01, -2.39291434e-01,
                -1.73326173e-01,  1.42412437e-01,  1.68169085e+00],
               [ 1.76693320e-01, -4.24251962e-01, -3.80376162e-01,
                -1.08277764e+00, -1.10939030e+00,  3.17721096e-01],
               [ 5.74302835e-01,  3.13225832e-01,  4.10135166e-01,
                 1.25819864e-01,  2.02728574e+00, -8.36395553e-01],
               [-8.22299859e-01,  3.66896163e-01, -3.72448587e-01,
                 1.43506039e-01,  2.78058492e-01,  6.40872897e-01],
               [-1.43546401e+00, -1.01993413e+00, -1.63479481e+00,
                -2.39549023e-01,  2.16124179e+00,  1.24044390e-01],
               [-3.08757088e-01,  5.64491405e-02,  4.54880298e-01,
                 3.34132625e-01, -3.39175998e-01,  3.26631289e-01],
               [-1.32316017e+00, -1.75577202e+00, -2.50142308e-01,
                -1.12842283e-01, -6.11034436e-01,  3.37970082e-01],
               [ 7.49216792e-02,  1.85929427e+00, -1.24200794e+00,
                 5.50207286e-01,  1.38738146e+00,  6.41696016e-01],
               [ 8.47209211e-01, -1.28471930e-01,  1.52034830e+00,
                 1.52685037e+00, -2.23480092e-01, -1.28662080e+00],
               [ 7.92943511e-01, -4.29441879e-01, -6.53815367e-02,
                 7.09751839e-01, -1.17067401e+00,  3.28341522e+00],
               [-9.90538959e-01, -1.25267317e+00, -4.06427559e-01,
                 6.29513454e-01, -2.02275495e-01,  7.29860220e-01],
               [-1.02956729e+00, -3.72288386e-01,  2.97354411e-01,
                 2.41555484e-01, -1.11517763e+00,  3.71989870e-01],
               [ 6.86378258e-01, -1.37917326e-01, -6.22020464e-01,
                -1.10976165e+00, -3.68346024e-01,  2.70505807e-01],
               [-3.90513830e-01,  3.04871939e-01,  8.51538275e-01,
                -4.20969496e-01, -5.86896815e-01,  8.02366824e-02],
               [-1.94302481e-01, -1.62109132e+00, -2.26622322e-01,
                 1.14690918e+00, -2.01053244e+00,  9.91035368e-01],
               [ 6.01294830e-01, -2.83633750e-01, -1.05105178e+00,
                 8.91662010e-01, -3.43522603e-01,  9.74126487e-01],
               [-1.08436820e+00,  7.23902420e-01, -1.42317643e+00,
                 1.05060841e-01, -2.64425047e-01, -6.62406263e-01],
               [-4.44448816e-01,  2.22848474e-01, -5.10270803e-01,
                 1.42185869e+00,  1.87171126e+00,  1.61013682e+00],
               [-7.53696908e-01,  7.77931973e-03, -3.59237899e-01,
                 1.25416465e+00,  5.33761870e-01, -7.80667917e-01],
               [-1.14529391e+00,  6.20250818e-01,  8.11937659e-01,
                -6.21590280e-01,  6.96670613e-02, -4.10056020e-01],
               [-6.14870740e-01, -4.33945137e-01,  6.76674938e-02,
                 4.15360904e-02, -1.13978311e-01,  3.68533533e-01],
               [ 6.06162127e-01, -7.18911365e-01,  8.66423769e-01,
                 5.87290991e-01, -2.23910955e-01,  4.34680540e-01],
               [ 9.74571643e-02, -5.47022109e-01, -8.96274654e-02,
                 1.39006131e+00,  9.98860504e-01,  3.08544608e-01],
               [-8.51270022e-01,  7.99954097e-01,  1.29578024e+00,
                 6.34025190e-01,  5.66979086e-01, -6.06127397e-01],
               [-5.80555301e-02,  8.81732112e-01,  8.86901862e-02,
                 3.99403825e-01,  1.50353577e+00, -1.77912312e-01],
               [-1.59908383e+00,  1.34080287e-01,  4.96945732e-01,
                 9.28604354e-01,  4.49296306e-01,  3.54456648e-01],
               [-7.44230974e-01, -9.76527150e-01, -8.60023183e-01,
                -7.18074065e-01, -6.58904985e-01,  1.71236903e+00],
               [ 2.04578099e+00, -1.28584348e-01,  4.40323690e-01,
                 3.07877358e-01,  1.86951269e-02,  1.48366922e+00],
               [-1.07658158e-02,  8.37810833e-01, -1.31414541e+00,
                -2.98097413e-01, -1.43182553e+00,  1.58386342e+00],
               [ 7.71481290e-01,  7.58576445e-01, -2.56555298e-01,
                -2.13460038e-01,  1.73296844e+00, -2.06244974e+00],
               [ 5.72230167e-02, -1.15266855e+00, -8.46896034e-01,
                -7.22457940e-01,  3.04546687e-02,  5.50732770e-01],
               [ 4.21865974e-01,  2.06643439e-01,  1.62730868e+00,
                 1.04339380e+00,  1.82554975e+00, -1.06906724e+00],
               [ 6.19743983e-01, -2.39606158e-01, -2.95052835e-01,
                -1.27847556e+00,  9.96971310e-01, -1.61134991e+00]]
        graph_gt = {'0': [('5', -3), ('2', -1)],
                     '1': [('2', -2)],
                     '2': [('3', -4), ('0', -1)],
                     '3': [('3', -4)],
                     '4': [('4', -2), ('0', -3)],
                     '5': [('3', -3), ('2', -4)]}
        return np.array(data), list(graph_gt.keys()),  graph_gt
    def test_base_time_series_model(self):
        
        data_array, var_names, graph_gt = self.get_data()

        data_train = data_array

        StandardizeTransform_ = StandardizeTransform()
        StandardizeTransform_.fit(data_train)

        data_train_trans = StandardizeTransform_.transform(data_train)

        data_train_obj = TimeSeriesData(data_train_trans, var_names=var_names)
        prior_knowledge = PriorKnowledge(forbidden_links={'0': ['1']}) # 1 cannot be a parent of 0

        model = BaseTimeSeriesAlgo(data=data_train_obj, prior_knowledge=prior_knowledge, use_multiprocessing=False)


        # check if get_all_parents() returns the correct set of parents
        all_parents = model.get_all_parents(target_var='0', max_lag=2)
        all_parents_gt = [('0', -1), ('0', -2), ('1', -1), ('1', -2),\
                          ('2', -1), ('2', -2), ('3', -1), ('3', -2),\
                          ('4', -1), ('4', -2), ('5', -1), ('5', -2)]
        self.assertTrue(set(all_parents)==set(all_parents_gt))

        # check if get_candidate_parents() returns the correct set of parents
        candidate_parents = model.get_candidate_parents(target_var='0', max_lag=2)
        candidate_parents_gt = [('0', -1), ('0', -2), ('2', -1), ('2', -2), ('3', -1), ('3', -2),\
                                 ('4', -1), ('4', -2), ('5', -1), ('5', -2)]
        self.assertTrue(set(candidate_parents)==set(candidate_parents_gt))

        # check if sort_parents() sorts the parents correctly based on pvalues
        parents_vals = {('A', -1): 0.1, ('A', -2): 0.3, ('B', -1): 0.2, ('B', -2): 0.01, ('C', -3): 0.05}
        sorted_parent = model.sort_parents(parents_vals)
        self.assertTrue(sorted_parent==[('A', -2), ('B', -1), ('A', -1), ('C', -3), ('B', -2)])


# if __name__ == "__main__":
#     unittest.main()
